#pragma once
//#include "probabilitymodel.h"
#include <stdexcept>
#include <vector>

class BitBuffer;
class ContextTreeNode;

class ContextTree
{
public:
	ContextTree(int blckSizeInBytes, int cntextSizeInBits);
	virtual ~ContextTree(void);

	double countTree(BitBuffer *bitSeq, int bitsToProcess);
	double countTree(BitBuffer *bitSeq);

	double appendBitToTree(char bit);
	double predict(char bit);

	void printTree();
	void printLUTTable();

private:
	double **LUTTable;
	int contextSizeInBits;
	unsigned int nodesCount;
	unsigned char actualBit;
	int blockSizeInBytes;
	int blockSizeInBits;
	int actualByte;
	std::vector<ContextTreeNode *> nodes;

	ContextTreeNode *root;

	unsigned char *initialContext;

	BitBuffer *bitSequence;
	BitBuffer *actualContext;
	BitBuffer *reversedContext;
	BitBuffer *recentSequence;
	BitBuffer *chosenBits;

	int actualContextSize;

	void generateChildren(ContextTreeNode *parent);
	int countNodesForTreeLevel(int level);

	bool checkIfActualContextExistsInSequence(int contextStartingIndex);

	void countZerosAndOnesOnNodes(ContextTreeNode *node);
	void updateZerosAndOnes(char bit);
	void rollbackLastBit(char bit);
	void countLUTTable();

	void countProbabilityOnNodes(ContextTreeNode *node);
	void updateProbabilitiesOnNodes();
	void updateBranch(ContextTreeNode *leaf);


public:
	class BadContextSize : public std::runtime_error
	{
	public:
		BadContextSize():std::runtime_error("Block size should be at least twice the size of context size!"){}
		~BadContextSize(){};
	};

};

